from PyCmpltrtok.common import sep
import torch

sep('4x5')
ten = torch.ones((4, 5), dtype=int)
ten[1, 3:5] = 0
ten[2, 2:5] = 0
ten[3, 1:5] = 0
ten_ = ten.clone()
print(ten.shape)
print(ten)

sep('view (4, 5, 1, 1)')
ten = ten.view(4, 5, 1, 1)
print(ten.shape)
print(ten)

sep('expand_as')
scores = torch.randn((4, 5, 3, 3), dtype=float)
ten = ten.expand_as(scores)
print(ten.shape)
print(ten)

sep('masked_select')
col = torch.Tensor(range(0, 40, 10)).view(4, 1)
row = torch.Tensor(range(5)).view(1, 5)
mat = col + row
print(mat.shape)
print(mat)
mat = mat.masked_select(ten_.bool())
print(mat.shape)
print(mat)
